DEFAULT_KB_CHANNEL = 11
DEFAULT_KB_DEVICE = None

from scapy.config import conf
setattr(conf, 'killerbee_channel', DEFAULT_KB_CHANNEL)
setattr(conf, 'killerbee_device', DEFAULT_KB_DEVICE)
setattr(conf, 'killerbee_nkey', None)
from scapy.base_classes import SetGen
from scapy.packet import Gen,Raw

import os
import time
from kbutils import randmac
import logging
log_killerbee = logging.getLogger('scapy.killerbee')

def __kb_send(kb, x, channel = None, inter = 0, loop = 0, count = None, verbose = None, realtime = None, *args, **kargs):
	if type(x) is str:
		x = Raw(load=x)
	if not isinstance(x, Gen):
		x = SetGen(x)
	if verbose is None:
		verbose = conf.verb

	n = 0
	if count is not None:
		loop = -count
	elif not loop:
		loop=-1
	dt0 = None
	try:
		while loop:
			for p in x:
				if realtime:
					ct = time.time()
					if dt0:
						st = dt0+p.time-ct
						if st > 0:
							time.sleep(st)
					else:
						dt0 = ct-p.time
				kb.inject(p.do_build()[:-2], channel = None, count = 1, delay = 0)	# [:-2] because the firmware adds the FCS
				n += 1
				if verbose:
					os.write(1,".")
				time.sleep(inter)
			if loop < 0:
				loop += 1
	except KeyboardInterrupt:
		pass
	return n

def __kb_recv(kb, count = 0, store = 1, prn = None, lfilter = None, stop_filter = None, verbose = None, timeout = None):
	kb.sniffer_on()
	if timeout is not None:
		stoptime = time.time()+timeout
	if verbose is None:
		verbose = conf.verb

	lst = []
	packetcount = 0
	remain = None
	while 1:
		try:
			if timeout is not None:
				remain = stoptime-time.time()
				if remain <= 0:
					break

			packet = kb.pnext()	# int(remain * 1000) to convert to seconds
			if packet == None: continue
			if verbose > 1:
				os.write(1, "*")
			packet = Dot15d4(packet[0])
			if lfilter and not lfilter(packet):
				continue
			packetcount += 1
			if store:
				lst.append(packet)
			if prn:
				r = prn(packet)
				if r is not None:
					print r
			if stop_filter and stop_filter(packet):
				break
			if count > 0 and packetcount >= count:
				break
		except KeyboardInterrupt:
			break

	kb.sniffer_off()
	kb.close()
	return lst

@conf.commands.register
def kbdev():
	"""List KillerBee recognized devices"""
	kbdev_info = kbutils.devlist()
	print 'Dev\tProduct String\tSerial Number\tCapabilities'
	for dev in kbdev_info:
		capabilities = []
		if dev[1] == 'KILLERB001':
			capabilities.append('SNIFF')
			capabilities.append('SETCHAN')
			capabilities.append('INJECT')
		elif dev[1] == 'RZUSBSTICK':
			capabilities.append('SNIFF')
			capabilities.append('SETCHAN')
		else:
			capabilities.append('UNKNOWN')

		print "%s\t%s\t%s\t%s" % (dev[0], dev[1], dev[2], ", ".join(capabilities))
	print ''
	print "Found %d devices." % len(kbdev_info)

@conf.commands.register
def kbsendp(pkt, channel = None, inter = 0, loop = 0, iface = None, verbose = None, realtime=None):
	"""Send a packet with KillerBee
channel:  802.15.4 channel to transmit/receive on
inter:    time to wait between tranmissions
loop:     number of times to process the packet list
iface:    KillerBee interface to use, or KillerBee() class instance
verbose:  set verbosity level
realtime: use packet's timestamp, bending time with realtime value"""
	if channel == None:
		channel = conf.killerbee_channel
	if not isinstance(iface, KillerBee):
		if iface is not None:
			kb = KillerBee(device = iface)
		else:
			kb = KillerBee(device = conf.killerbee_device)
		kb.set_channel(channel)
	else:
		kb = iface
	
	pkts_out = __kb_send(kb, pkt, inter = inter, loop = loop, count = None, verbose = verbose, realtime = realtime)
	print "\nSent %i packets." % pkts_out

@conf.commands.register
def kbsrp(pkt, channel = None, inter = 0, count = 0, iface = None, store = 1, prn = None, lfilter = None, timeout = None, verbose = None, realtime = None):
	"""Send and receive packets with KillerBee
channel:  802.15.4 channel to transmit/receive on
inter:    time to wait between tranmissions
count:    number of packets to capture. 0 means infinity
iface:    KillerBee interface to use, or KillerBee() class instance
store:    wether to store sniffed packets or discard them
prn:      function to apply to each packet. If something is returned,
          it is displayed. Ex:
          ex: prn = lambda x: x.summary()
lfilter:  python function applied to each packet to determine
          if further action may be done
          ex: lfilter = lambda x: x.haslayer(Padding)
timeout:  stop sniffing after a given time (default: None)
verbose:  set verbosity level
realtime: use packet's timestamp, bending time with realtime value"""
	if verbose is None:
		verbose = conf.verb
	if channel == None:
		channel = conf.killerbee_channel
	if not isinstance(iface, KillerBee):
		if iface is not None:
			kb = KillerBee(device = iface)
		else:
			kb = KillerBee(device = conf.killerbee_device)
		kb.set_channel(channel)
	else:
		kb = iface
	
	pkts_out = __kb_send(kb, pkt, inter = inter, loop = 0, count = None, verbose = verbose, realtime = realtime)
	if verbose:
		print "\nSent %i packets." % pkts_out

	pkts_in = __kb_recv(kb, count = count, store = store, prn = prn, lfilter = lfilter, verbose = verbose, timeout = timeout)
	if verbose:
		print "\nReceived %i packets." % len(pkts_in)
	return plist.PacketList(pkts_in, 'Results')

@conf.commands.register
def kbsrp1(pkt, channel = None, inter = 0, iface = None, store = 1, prn = None, lfilter = None, timeout = None, verbose = None, realtime = None):
	"""Send and receive packets with KillerBee and return only the first answer"""
	return kbsrp(pkt, channel = channel, inter = inter, count = 1, iface = iface, store = store, prn = prn, lfilter = lfilter, timeout = timeout, verbose = verbose, realtime = realtime)

@conf.commands.register	
def kbsniff(channel = None, count = 0, iface = None, store = 1, prn = None, lfilter = None, stop_filter = None, verbose = None, timeout = None):
	"""Sniff packets with KillerBee
channel:  802.15.4 channel to transmit/receive on
count:    number of packets to capture. 0 means infinity
iface:    KillerBee interface to use, or KillerBee() class instance
store:    wether to store sniffed packets or discard them
prn:      function to apply to each packet. If something is returned,
          it is displayed. Ex:
          ex: prn = lambda x: x.summary()
lfilter:  python function applied to each packet to determine
          if further action may be done
          ex: lfilter = lambda x: x.haslayer(Padding)
timeout:  stop sniffing after a given time (default: None)"""
	if channel == None:
		channel = conf.killerbee_channel
	if not isinstance(iface, KillerBee):
		if iface is not None:
			kb = KillerBee(device = iface)
		else:
			kb = KillerBee(device = conf.killerbee_device)
		kb.set_channel(channel)
	else:
		kb = iface
	return plist.PacketList(__kb_recv(kb, count = count, store = store, prn = prn, lfilter = lfilter, stop_filter = stop_filter, verbose = verbose, timeout = timeout), 'Sniffed')

@conf.commands.register
def kbrdpcap(filename, count = -1, skip = 0):
	"""Read a pcap file with the KillerBee library
Wraps the PcapReader to return scapy packet object from pcap files.
This uses the killerbee internal methods instead of the scapy native methods."""
	cap = PcapReader(filename)
	lst = []
	packetcount = 0
	if count > 0:
		count += skip
	
	while 1:
		packet = cap.pnext()
		packetcount += 1
		if packet[1] == None:
			break
		if skip > 0 and packetcount <= skip:
			continue
		packet = Dot15d4(packet[1])
		lst.append(packet)
		if count > 0 and packetcount >= count:
			break
	return plist.PacketList(lst, os.path.basename(filename))

@conf.commands.register
def kbwrpcap(save_file, pkts):
	"""Write a pcap using the KillerBee library"""
	pd = PcapDumper(DLT_IEEE802_15_4, save_file, ppi=False)
	for packet in pkts:
		pd.pcap_dump(str(packet))
	pd.close()

@conf.commands.register
def kbrddain(filename, count = -1, skip = 0):
	"""Read a dain tree file with the KillerBee library
Wraps the DainTreeReader to return scapy packet object from dain tree files."""
	cap = DainTreeReader(filename)
	lst = []
	packetcount = 0
	
	while 1:
		packet = cap.pnext()
		packetcount += 1
		if packet[1] == None:
			break
		if skip > 0 and packetcount <= skip:
			continue
		packet = Dot15d4(packet[1])
		lst.append(packet)
		if count > 0 and packetcount >= count:
			break
	return plist.PacketList(lst, os.path.basename(filename))

@conf.commands.register
def kbwrdain(save_file, pkts):
	"""Write a dain tree file using the KillerBee library"""
	dt = DainTreeDumper(save_file)
	for packet in pkts:
		dt.pwrite(str(packet))
	dt.close()

@conf.commands.register
def kbkeysearch(packet, searchdata, ispath = True, skipfcs = True, raw = False):
	"""Search a binary file for the encryption key to an encrypted packet"""
	if 'fcf_security' in packet.fields and packet.fcf_security == 0:
		raise Exception('Packet Not Encrypted (fcf_security Not Set)')
	if ispath:
		searchdata = open(searchdata, 'r').read()
	packet = packet.do_build()
	if skipfcs:
		packet = packet[:-2]
	offset = 0
	keybytes = []
	d = Dot154PacketParser()
	searchdatalen = len(searchdata)
	while (offset < (searchdatalen - 16)):
		if d.decrypt(packet, searchdata[offset:offset+16]) != '':
			if raw:
				return ''.join(searchdata[offset + i] for i in xrange(0, 16))
			else:
				return ':'.join("%02x" % ord(searchdata[offset + i]) for i in xrange(0, 16))
		else:
			offset+=1
	return None

@conf.commands.register
def kbgetnetworkkey(pkts):
	"""Search packets for a plaintext key exchange
returns the first one found"""
	if not isinstance(pkts, Gen):
		pkts = SetGen(pkts)
	for packet in pkts:
		packet = str(packet)
		zmac = Dot154PacketParser()
		znwk = ZigBeeNWKPacketParser()
		zaps = ZigBeeAPSPacketParser()
		try:
			# Process MAC layer details
			zmacpayload = zmac.pktchop(packet)[-1]
			if zmacpayload == None:
				continue

			# Process NWK layer details
			znwkpayload = znwk.pktchop(zmacpayload)[-1]
			if znwkpayload == None:
				continue

			# Process the APS layer details
			zapschop = zaps.pktchop(znwkpayload)
			if zapschop == None:
				continue

			# See if this is an APS Command frame
			apsfc = ord(zapschop[0])
			if (apsfc & ZBEE_APS_FCF_FRAME_TYPE) != ZBEE_APS_FCF_CMD:
				continue

			# Delivery mode is Normal Delivery (0)
			apsdeliverymode = (apsfc & ZBEE_APS_FCF_DELIVERY_MODE) >> 2
			if apsdeliverymode != 0:
				continue

			# Ensure Security is Disabled
			if (apsfc & ZBEE_APS_FCF_SECURITY) == 1:
				continue

			zapspayload = zapschop[-1]

			# Check payload length, must be at least 35 bytes
			# APS cmd | key type | key | sequence number | dest addr | src addr
			if len(zapspayload) < 35:
				continue

			# Check for APS command identifier Transport Key (0x05)
			if ord(zapspayload[0]) != 5:
				continue

			# Transport Key Frame, get the key type.  Network Key is 0x01, no
			# other keys should be sent in plaintext
			if ord(zapspayload[1]) != 1:
				continue

			# Reverse these fields
			networkkey = zapspayload[2:18][::-1]
			destaddr = zapspayload[19:27][::-1]
			srcaddr = zapspayload[27:35][::-1]

			key_bytes = []
			dst_mac_bytes = []
			src_mac_bytes = []
			key = {}
			key['key'] = ':'.join("%02x" % ord(networkkey[x]) for x in xrange(16))
			key['dst'] = ':'.join("%02x" % ord(destaddr[x]) for x in xrange(8))
			key['src'] = ':'.join("%02x" % ord(srcaddr[x]) for x in xrange(8))
			return key
		except:
			continue
	return { }

@conf.commands.register
def kbtshark(store = 0, *args,**kargs):
	"""Sniff packets using KillerBee and print them calling pkt.show()"""
	return kbsniff(prn=lambda x: x.display(), store = store, *args, **kargs)

@conf.commands.register
def kbrandmac(length = 8):
	"""Returns a random MAC address using a list valid OUI's from ZigBee device manufacturers."""
	return randmac(length)

@conf.commands.register
def kbdecrypt(pkt, key = None, verbose = None):
	"""Decrypt Zigbee frames using AES CCM* with 32-bit MIC"""
	if verbose is None:
		verbose = conf.verb
	if key == None:
		if conf.killerbee_nkey == None:
			log_killerbee.error("Cannot find decryption key. (Set conf.killerbee_nkey)")
			return None
		key = conf.killerbee_nkey
	if len(key) != 16:
		log_killerbee.error("Invalid decryption key, must be a 16 byte string.")
		return None
	elif not pkt.haslayer(ZigbeeSecurityHeader) or not pkt.haslayer(ZigbeeNWK):
		log_killerbee.error("Cannot decrypt frame without a ZigbeeSecurityHeader.")
		return None
	try:
		import zigbee_crypt
	except ImportError:
		log_killerbee.error("Could not import zigbee_crypt extension, cryptographic functionality is not available.")
		return None

	from struct import pack
	f = pkt.getlayer(ZigbeeSecurityHeader).fields

	mic = pack(">I", f['mic'])
	encrypted = f['data']

	nonce = ""	# build the nonce
	nonce += pack(">Q", f['source'])
	nonce += pack(">I", f['fc'])
	fc = (f['reserved1'] << 6) | (f['extended_nonce'] << 5) | (f['key_type'] << 3) | f['reserved2']
	nonce += chr(fc | 0x05)

	if verbose > 2:
		print "Decrypt Details:"
		print "\tKey:            " + key.encode('hex')
		print "\tNonce:          " + nonce.encode('hex')
		print "\tMic:            " + mic.encode('hex')
		print "\tEncrypted Data: " + encrypted.encode('hex')
		
	crop_size = 4 + 2 + len(pkt.getlayer(ZigbeeSecurityHeader).fields['data'])	# the size of all the zigbee crap, minus the length of the encrypted data, mic and FCS
	
	# the Security Control Field flags have to be adjusted before this is calculated, so we store their original values so we can reset them later
	reserved2 = pkt.getlayer(ZigbeeSecurityHeader).fields['reserved2']
	pkt.getlayer(ZigbeeSecurityHeader).fields['reserved2'] = (pkt.getlayer(ZigbeeSecurityHeader).fields['reserved2'] | 0x05)
	zigbeeData = pkt.getlayer(ZigbeeNWK).do_build()
	zigbeeData = zigbeeData[:-crop_size]
	pkt.getlayer(ZigbeeSecurityHeader).fields['reserved2'] = reserved2
	
	(payload, micCheck) = zigbee_crypt.decrypt_ccm(key, nonce, mic, encrypted, zigbeeData)

	if verbose > 2:
		print "\tDecrypted Data: " + payload.encode('hex')

	frametype = pkt.getlayer(ZigbeeNWK).fields['frametype']
	if frametype == 0:
		payload = ZigbeeAppDataPayload(payload)
	elif frametype == 1:
		payload = ZigbeeNWKCommandPayload(payload)
	else:
		payload = Raw(payload)
	if micCheck == 1:
		return (payload, True)
	else:
		return (payload, False)

@conf.commands.register
def kbencrypt(pkt, data, key = None, verbose = None):
	"""Encrypt Zigbee frames using AES CCM* with 32-bit MIC"""
	if verbose is None:
		verbose = conf.verb
	if key == None:
		if conf.killerbee_nkey == None:
			log_killerbee.error("Cannot find decryption key. (Set conf.killerbee_nkey)")
			return None
		key = conf.killerbee_nkey
	if len(key) != 16:
		log_killerbee.error("Invalid encryption key, must be a 16 byte string.")
		return None
	elif not pkt.haslayer(ZigbeeSecurityHeader) or not pkt.haslayer(ZigbeeNWK):
		log_killerbee.error("Cannot encrypt frame without a ZigbeeSecurityHeader.")
		return None
	try:
		import zigbee_crypt
	except ImportError:
		log_killerbee.error("Could not import zigbee_crypt extension, cryptographic functionality is not available.")
		return None
		
	from struct import unpack
	f = pkt.getlayer(ZigbeeSecurityHeader).fields

	f['data'] = ''	# explicitly clear it out, this should go without say
	
	if isinstance(data, Packet):
		decrypted = data.do_build()
	else:
		decrypted = data

	nonce = ""	# build the nonce
	nonce += pack(">Q", f['source'])
	nonce += pack(">I", f['fc'])
	fc = (f['reserved1'] << 6) | (f['extended_nonce'] << 5) | (f['key_type'] << 3) | f['reserved2']
	nonce += chr(fc | 0x05)

	if verbose > 2:
		print "Encrypt Details:"
		print "\tKey:            " + key.encode('hex')
		print "\tNonce:          " + nonce.encode('hex')
		print "\tDecrypted Data: " + decrypted.encode('hex')
		
	crop_size = 4 + 2 # the size of all the zigbee crap, minus the length of the mic and FCS
	
	# the Security Control Field flags have to be adjusted before this is calculated, so we store their original values so we can reset them later
	reserved2 = pkt.getlayer(ZigbeeSecurityHeader).fields['reserved2']
	pkt.getlayer(ZigbeeSecurityHeader).fields['reserved2'] = (pkt.getlayer(ZigbeeSecurityHeader).fields['reserved2'] | 0x05)
	zigbeeData = pkt.getlayer(ZigbeeNWK).do_build()
	zigbeeData = zigbeeData[:-crop_size]
	pkt.getlayer(ZigbeeSecurityHeader).fields['reserved2'] = reserved2
	
	(payload, mic) = zigbee_crypt.encrypt_ccm(key, nonce, 4, decrypted, zigbeeData)

	if verbose > 2:
		print "\tEncrypted Data: " + payload.encode('hex')
		print "\tMic:            " + mic.encode('hex')
	
	# Set pkt's values to reflect the encrypted ones to it's ready to be sent
	f['data'] = payload
	f['mic'] = unpack(">I", mic)[0]
	return pkt
